#include "poll.h"

static int compare_node(const struct rbnode* r1, const struct rbnode* r2)
{
	struct timeo_node *node1, *node2;

	node1 = rbnode_entry(r1, struct timeo_node, node);
	node2 = rbnode_entry(r2, struct timeo_node, node);
	return (int)(node1->expire - node2->expire);
}

static inline void timeo_node_delete(struct rbtree* rbtree, struct timeo_node* node)
{
	if(!rbnode_empty(&node->node))
		rbtree_delete(rbtree, &node->node);
}

static inline void timeo_node_active(struct rbtree* rbtree, struct timeo_node* node)
{
	if(node->seconds > 0) {
		node->expire = ae_timestamp(NULL) + node->seconds * 1000;
		rbtree_insert(rbtree, &node->node);
	}
}

static inline void timeo_node_change(struct rbtree* rbtree, struct timeo_node* node)
{
	timeo_node_delete(rbtree, node);
	timeo_node_active(rbtree, node);
}

static inline struct timeo_node* pickup_expire_node(struct rbtree* rbtree, ae_time_t current, int* milliseconds)
{
	struct timeo_node* node;
	struct rbnode* rbnode;

	if(rbtree_empty(rbtree)) {
		*milliseconds = 0;
		return NULL;
	}
	rbnode = rbtree_first(rbtree);
	node = rbnode_entry(rbnode, struct timeo_node, node);
	if(node->expire <= current) {
		node->expire = current + node->seconds * 1000;
		rbtree_update(rbtree, rbnode);
		return node;
	}
	*milliseconds = (int)(node->expire - current);
	return NULL;
}

static void pickup_expire_customer(struct poll_provider* provider, struct poll_customer **customer, int *events, int* milliseconds)
{
	ae_time_t current;
	struct timeo_node* node;
	int write_treeut = 0, read_treeut = 0;

	*customer = NULL;
	*milliseconds = 0;

	ae_timestamp(&current);
	mutex_lock(&provider->write_lock);
	node = pickup_expire_node(&provider->write_tree, current, &write_treeut);
	mutex_unlock(&provider->write_lock);
	if(node) {
		*customer = container_of(node, struct poll_customer, write_node);
		*events = AE_POLL_SNDTIMEO;
	} else {
		mutex_lock(&provider->read_lock);
		node = pickup_expire_node(&provider->read_tree, current, &read_treeut);
		mutex_unlock(&provider->read_lock);
		if(node) {
			*customer = container_of(node, struct poll_customer, read_node);
			*events = AE_POLL_RCVTIMEO;
		}
		else {
			if(read_treeut > 0 && write_treeut > 0)
				*milliseconds = AE_MIN(read_treeut, write_treeut);
			else if(read_treeut > 0)
				*milliseconds = read_treeut;
			else if(write_treeut > 0)
				*milliseconds = write_treeut;
		}
	}
}

static void customer_detach(struct poll_customer* customer)
{
	struct poll_provider* provider = customer->__provider;
	atomic_dec(&provider->nr_customer);
}

static void customer_reset_timeout(struct poll_customer* customer, int type)
{
	struct poll_provider* provider = customer->__provider;
	if(type == AE_POLL_READ) {
		mutex_lock(&provider->read_lock);
		timeo_node_change(&provider->read_tree, &customer->read_node);
		mutex_unlock(&provider->read_lock);
	}
	else {
		mutex_lock(&provider->write_lock);
		timeo_node_change(&provider->write_tree, &customer->write_node);
		mutex_unlock(&provider->write_lock);
	}
}

static void customer_watch_event(struct poll_customer* customer, int type)
{
	struct poll_provider* provider = customer->__provider;
	if(type == AE_POLL_READ) {
		mutex_lock(&provider->read_lock);
		timeo_node_active(&provider->read_tree, &customer->read_node);
		mutex_unlock(&provider->read_lock);
	}
	else {
		mutex_lock(&provider->write_lock);
		timeo_node_active(&provider->write_tree, &customer->write_node);
		mutex_unlock(&provider->write_lock);
	}
}

static void customer_event_arrayed(struct poll_customer* customer, int type)
{
	struct poll_provider* provider = customer->__provider;
	if(type == AE_POLL_READ) {
		mutex_lock(&provider->read_lock);
		timeo_node_delete(&provider->read_tree, &customer->read_node);
		mutex_unlock(&provider->read_lock);
	}
	else {
		mutex_lock(&provider->write_lock);
		timeo_node_delete(&provider->write_tree, &customer->write_node);
		mutex_unlock(&provider->read_lock);
	}
	customer->__method->event_arrayed(customer, type);
}

static int customer_maintain(struct poll_customer* customer, int is_down)
{
	return customer->__method->maintain(customer, is_down);
}

static struct poll_provider_method poll_method = {
	customer_detach,
	customer_watch_event,
	customer_reset_timeout
};

int ae_poll_create(ae_object_t* poll)
{
	struct poll_provider* provider;
	poll_engine* rpoll = (poll_engine*)ae_mallocz(sizeof(poll_engine));
	if(rpoll == NULL)
		return -1;

#if defined(AE_WIN32)
	rpoll->poll = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
#elif defined(AE_LINUX)
	rpoll->poll = epoll_create(1024);
#elif defined(AE_FREEBSD)
	rpoll->poll = kqueue();
#endif
	if(rpoll->poll == INVALID_POLL) {
		ae_free(rpoll);
		return -1;
	}

	provider = TO_POLL_PROVIDER(rpoll);
	provider->__method = &poll_method;
	mutex_init(&provider->read_lock);
	mutex_init(&provider->write_lock);
	rbtree_init(&provider->read_tree, compare_node);
	rbtree_init(&provider->write_tree, compare_node);
	SET_OBJECT_TYPE(rpoll, OBJ_EVENT);
	*poll = (ae_object_t) rpoll;
	return 0;
}

int ae_poll_destory(ae_object_t poll)
{
	poll_engine* rpoll = (poll_engine*)poll;
	struct poll_provider* provider = TO_POLL_PROVIDER(rpoll);

	if(!CHECK_OBJECT(rpoll, OBJ_EVENT) || atomic_read(&provider->nr_customer) != 0)
		return -1;

	closepoll(rpoll->poll);
	mutex_destory(&provider->read_lock);
	mutex_destory(&provider->write_lock);
	ae_free(rpoll);
	return 0;
}

int ae_poll_attach(ae_object_t poll, ae_object_t object, void* context)
{
	int retval = -1, events;
	__poll_t handle;
#if defined(AE_LINUX)
	struct epoll_event edata;
#elif defined(AE_FREEBSD)
	struct kevent kev[2];
#endif
	poll_engine* rpoll = (poll_engine*)poll;
	struct poll_provider* provider = TO_POLL_PROVIDER(rpoll);
	struct poll_customer* customer = TO_POLL_CUSTOMER(object);

	if(!CHECK_OBJECT(rpoll, OBJ_EVENT) || !CHECK_OBJECT_CANEVENT(customer) || customer->__provider != NULL)
		return retval;

	customer->context = context;
	customer->__provider = provider;
	if(customer->__method->before_attath(customer, &handle, &events) != 0) {
		customer->context = NULL;
		customer->__provider = NULL;
		return retval;
	}

#if defined(AE_WIN32)
	if(CreateIoCompletionPort(handle, rpoll->poll, (ULONG_PTR)customer, 0) != NULL)
		retval = 0;
#elif defined(AE_LINUX)
	edata.events = EPOLLIN | EPOLLET;
	if(events & AE_POLL_WRITE)
		edata.events |= EPOLLOUT;
	if(events & AE_POLL_HANGUP)
		edata.events |= EPOLLRDHUP | EPOLLHUP;
	edata.data.ptr = customer;
	if(epoll_ctl(rpoll->poll, EPOLL_CTL_ADD, handle, &edata) == 0)
		retval = 0;
#elif defined(AE_FREEBSD)
	EV_SET(&kev[0], handle, EVFILT_READ, EV_ADD | EV_ENABLE | EV_CLEAR, 0, 0, customer);
	EV_SET(&kev[1], handle, EVFILT_WRITE, EV_ADD | EV_ENABLE | EV_CLEAR, 0, 0, customer);
	if(kevent(rpoll->poll, kev, 2, NULL, 0, NULL) == 0)
		retval = 0;
#endif
	if(retval == 0) {
		customer->__method->after_attath(customer, 0);
		atomic_inc(&provider->nr_customer);
	}
	else {
		customer->__method->after_attath(customer, 1);
		customer->context = NULL;
		customer->__provider = NULL;
	}
	return retval;
}

#if defined(AE_WIN32)
static int inline_always poll_wait(poll_engine* rpoll, struct poll_customer **customer, int *events, int milliseconds)
{
	int is_netdown;
	DWORD bytes, timeout, last_error;
	struct poll_customer *_customer;
	IOCPOVERLAPPED *overlapped = NULL;

	if(milliseconds < 0)
		timeout = INFINITE;
	else
		timeout = milliseconds;

	do {
		*customer = NULL;
		*events = 0;
		is_netdown = 0;
		if(!GetQueuedCompletionStatus(rpoll->poll, &bytes, (PULONG_PTR)&_customer, (LPOVERLAPPED *)&overlapped, timeout)) {
			last_error = GetLastError();
			if(last_error == WAIT_TIMEOUT)
				return 0;
			if(_customer == NULL || overlapped == NULL)
				return -1;
			if(last_error == ERROR_NETNAME_DELETED)
				is_netdown = 1;
		}
		if(overlapped->type == AE_POLL_READ)
			customer_event_arrayed(_customer, AE_POLL_READ);
		else if(overlapped->type == AE_POLL_WRITE)
			customer_event_arrayed(_customer, AE_POLL_WRITE);
	} while(customer_maintain(_customer, is_netdown) == 0);

	*customer = _customer;
	*events = overlapped->type;
	if(is_netdown)
		*events |= AE_POLL_HANGUP;
	return 1;
}
#elif defined(AE_LINUX)
static int inline_always poll_wait(poll_engine* rpoll, struct poll_customer **customer, int *events, int milliseconds)
{
	int ready, is_netdown;
	sigset_t origmask;
	struct epoll_event edata;
	struct poll_customer *_customer = NULL;

	if(milliseconds < 0)
		milliseconds = -1;

	do {
		*customer = NULL;
		*events = 0;
		is_netdown = 0;
		ready = epoll_pwait(rpoll->poll, &edata, 1, milliseconds, &origmask);
		if(ready == 0)
			return 0;
		else if(ready < 0) {
			if(errno != EINTR)
				return -1;
			else
				continue;
		}
		if ((edata.events & EPOLLHUP) || (edata.events & EPOLLRDHUP) || (edata.events & EPOLLERR))
			is_netdown = 1;
		_customer = (struct poll_customer*) edata.data.ptr;
		if(edata.events & EPOLLIN)
			customer_event_arrayed(_customer, AE_POLL_READ);
		if(edata.events & EPOLLOUT)
			customer_event_arrayed(_customer, AE_POLL_WRITE);
	} while(customer_maintain(_customer, is_netdown) == 0);

	*customer = _customer;
	if(edata.events & EPOLLIN)
		*events |= AE_POLL_READ;
	if(edata.events & EPOLLOUT)
		*events |= AE_POLL_WRITE;
	if((edata.events & EPOLLHUP) || (edata.events & EPOLLRDHUP))
		*events |= AE_POLL_HANGUP;
	if(edata.events & EPOLLERR)
		*events |= AE_POLL_ERROR;
	return 1;
}
#elif defined(AE_FREEBSD)
static int inline_always poll_wait(poll_engine* rpoll, struct poll_customer **customer, int *events, int milliseconds)
{
	int ready, is_netdown;
	struct kevent edata, kev[1];
	struct timespec *timeout, tsdata;
	struct poll_customer *_customer = NULL;

	if(milliseconds < 0)
		timeout = NULL;
	else
	{
		timeout = &tsdata;
		timeout->tv_sec = milliseconds / 1000;
		timeout->tv_nsec = (milliseconds % 1000) * 1000;
	}

	do {
		*customer = NULL;
		*events = 0;
		is_netdown = 0;
		ready = kevent(rpoll->poll, NULL, 0, &edata, 1, timeout);
		if(ready == 0)
			return 0;
		else if(ready < 0) {
			if(errno != EINTR)
				return -1;
			else
				continue;
		}
		_customer = (struct poll_customer*) edata.udata;
		if(edata.flags & (EV_EOF | EV_ERROR)) {
			is_netdown = 1;
			EV_SET(&kev[0], edata.ident, edata.filter, EV_DELETE, 0, 0, NULL);
            kevent(rpoll->poll, kev, 1, NULL, 0, NULL);
		}
		if(edata.filter == EVFILT_READ)
			customer_event_arrayed(_customer, AE_POLL_READ);
		else if(edata.filter == EVFILT_WRITE)
			customer_event_arrayed(_customer, AE_POLL_WRITE);
	} while(customer_maintain(_customer, is_netdown) == 0);

	*customer = _customer;
	if(edata.filter == EVFILT_READ)
		*events |= AE_POLL_READ;
	if(edata.filter == EVFILT_WRITE)
		*events |= AE_POLL_WRITE;
	if(edata.flags & EV_EOF)
		*events |= AE_POLL_HANGUP;
	if(edata.flags & EV_ERROR)
		*events |= AE_POLL_ERROR;
	return 1;
}
#endif

int ae_poll_wait(ae_object_t poll, ae_object_t *object, int *events, void** context, int milliseconds)
{
	int retval, near_expire;
	struct poll_customer *customer;
	poll_engine* rpoll = (poll_engine*)poll;

	if(!CHECK_OBJECT(rpoll, OBJ_EVENT) || object == NULL || events == NULL || context == NULL)
		return -1;

	pickup_expire_customer(TO_POLL_PROVIDER(rpoll), &customer, events, &near_expire);
	if(customer != NULL) {
		*object = (ae_object_t)TO_OBJECT(customer);
		*context = customer->context;
		return 1;
	}
	else if(near_expire > 0) {
		if(milliseconds < 0)
			milliseconds = near_expire;
		else
			milliseconds = AE_MIN(near_expire, milliseconds);
	}

	retval = poll_wait(rpoll, &customer, events, milliseconds);
	if(retval == 1) {
		*object = (ae_object_t)TO_OBJECT(customer);
		*context = customer->context;
	}
	else if(retval == 0) {
		pickup_expire_customer(TO_POLL_PROVIDER(rpoll), &customer, events,  &near_expire);
		if(customer != NULL) {
			*object = (ae_object_t)TO_OBJECT(customer);
			*context = customer->context;
			retval = 1;
		}
	}
	return retval;
}
